import glob, sys
import numpy as np
import os
sys.dont_write_bytecode = True

"""-----------------------------Split the data-------------------------------"""
def split_data_2_txt(datapath,savepath,rate,suffix, which='training', name=None):
    imgs = glob.glob(datapath+"/*"+suffix)
    print('The number of images is: %d'%len(imgs))

    randidx = np.arange(len(imgs))
    np.random.shuffle(randidx)

    # distribute the data to the train, test, and val path
    trainnum = int(len(imgs) * rate[0])
    valnum = int(len(imgs) * rate[1])

    traintxt = savepath + os.sep + ('train.txt' if not name and (which == 'training') else name)
    valtxt = savepath + os.sep + ('val.txt' if not name and (which == 'validation') else name)
    testtxt = savepath + os.sep + ('test.txt' if not name and (which == 'test') else name)
    alltxt = savepath + os.sep + 'all.txt'

    trainimgs = [imgs[idx] for idx in randidx[:trainnum]]
    valimgs = [imgs[idx] for idx in randidx[trainnum:trainnum+valnum]]
    testimgs = [imgs[idx] for idx in randidx[trainnum+valnum:]]

    if which == 'training':
        write_path_2_txt(traintxt,trainimgs)
    elif which == 'validation':
        write_path_2_txt(valtxt,valimgs)
    elif which == 'test':
        write_path_2_txt(testtxt,testimgs)
    # write_path_2_txt(alltxt,imgs)

def write_path_2_txt(txtpath,imgpaths):
    fidf = open(txtpath,'w')
    num = len(imgpaths)
    for idx in range(num):
        imgname = imgpaths[idx]
        fidf.write(imgname)
        if idx!=(num-1):
            fidf.write("\n")
    fidf.close()

if __name__ == '__main__':
    # train_path = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask/img'
    # val_path = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask/img'

    # train_path = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask/img'
    # val_path = '/home/yusongli/_dataset/shidaoai/img/_out/wangqifeng_spacial_cropped_96_224_224_img_mask/img'

    savepath = 'task/01_seg/2D-Unet/_conf'

    # split_data_2_txt(data_path, savepath, rate=[0.8, 0.2], suffix='.nii.gz')
    split_data_2_txt(train_path, savepath, rate=[1, 0], suffix='.nii.gz')
    split_data_2_txt(val_path, savepath, rate=[0, 1], suffix='.nii.gz')
